# -*- coding: utf-8 -*-
"""
@Time    : 2022/9/16 09:30
@Author  : huanghuang
"""
import os
from typing import List, Any

import openpyxl
from openpyxl import Workbook
from openpyxl.worksheet.worksheet import Worksheet

from common.pkg.str_id import get_short_uuid
from common.g.g import G_SETTINGS


def generate_excel(header: List[str], data: List[List[Any]], filename: str = None, sheet_name: str = None) -> str:
    """
        根据给出来的数据生成excel文件
    :param header: excel的header
    :param data: 数据
    :param filename: 文件名，如果没有将自动生成
    :param sheet_name:
    :return: str
    """
    filename = f"{G_SETTINGS.static_root}/{filename}.xlsx" or f"{G_SETTINGS.static_root}/{get_short_uuid()}.xlsx"
    wb = openpyxl.Workbook()
    # 激活第一个sheet
    ws = wb.active  # type: Worksheet
    if sheet_name:
        ws.title = sheet_name
    # 生成header
    ws.append(header)
    # 生成数据
    if data:
        for each in data:
            ws.append(each)
    # 保存
    wb.save(filename)
    wb.close()
    return filename


def read_data_from_excel(data: bytes, use_name_or_index: str, remove: bool = True):
    """
        读取excel数据并返回
    :param data: excel的二进制数据
    :param use_name_or_index: 使用excel的header中的名字还是下标
    :param remove: 使用完之后是否删除文件
    :return:
    """
    filename = f"{G_SETTINGS.static_root}/{get_short_uuid()}.xlsx"
    with open(filename, "wb") as f:
        f.write(data)
    wb = openpyxl.load_workbook(filename, read_only=True)  # type: Workbook
    ws = wb.active  # type: Worksheet
    data_start_row = 2  # 第一行是header， 第二行才是数据行
    header_col = ws[1]  # 获取第一行的header
    total_rows = ws.max_row  # 获取所有的行
    data = []
    while data_start_row <= total_rows:
        _d = {}
        for i, d in enumerate(header_col):
            if use_name_or_index == "name":
                _d[header_col[i].value] = ws.cell(data_start_row, i+1).value
            else:
                _d[str(i)] = ws.cell(data_start_row, i + 1).value
        data.append(_d)
        data_start_row += 1
    if remove:
        os.remove(filename)  # 清除数据
    return data


if __name__ == '__main__':
    import datetime
    header = ["姓名", "年龄", "生日"]
    data = [
        ["huang", 12, datetime.datetime.now()],
        ["yue", 14],
        ["ling", 14]
    ]
    filename = generate_excel(header, data, sheet_name="学生表")
